#!/usr/bin/env python3

# Python script to run and analyse MMS test

from __future__ import division
from __future__ import print_function
try:
  from builtins import str
except:
  pass

from boututils.run_wrapper import shell, shell_safe, launch_safe
from boutdata.collect import collect

from numpy import sqrt, max, abs, mean, array, log



print("Making MMS diffusion test")
shell_safe("make > make.log")

# List of NX values to use
nxlist = [4, 8, 16, 32, 64, 128]

nout = 1
timestep = 0.1

nproc = 1

error_2   = []  # The L2 error (RMS)
error_inf = []  # The maximum error

for nx in nxlist:
    args = "mesh:nx="+str(nx)+" nout="+str(nout)+" timestep="+str(timestep)
    
    print("Running with " + args)

    # Delete old data
    shell("rm data/BOUT.dmp.*.nc")
    
    # Command to run
    cmd = "./cyto "+args
    # Launch using MPI
    s, out = launch_safe(cmd, nproc=nproc, pipe=True)

    # Save output to log file
    f = open("run.log."+str(nx), "w")
    f.write(out)
    f.close()

    # Collect data
    E_N = collect("E_N", tind=[nout,nout], path="data", info=False)

    E_N = E_N[0,:,0,0]

    # Average error over domain, not including guard cells
    l2 = sqrt(mean(E_N[1:-1]**2))
    linf = max(abs( E_N[1:-1] ))
    
    error_2.append( l2 )
    error_inf.append( linf )

    print("Error norm: l-2 %f l-inf %f" % (l2, linf))

# Calculate grid spacing
dx = 1. / (array(nxlist) - 2.)

order = log(error_2[-1] / error_2[-2]) / log(dx[-1] / dx[-2])
print("Convergence order = %f" % (order))

# Attempt to plot errors
try:
  import matplotlib.pyplot as plt

  plt.plot(dx, error_2, '-o', label=r'$l^2$')
  plt.plot(dx, error_inf, '-x', label=r'$l^\infty$')
  plt.plot(dx, error_2[-1]*(dx/dx[-1])**order, '--', label="Order %.1f"%(order))

  plt.legend(loc="upper left")
  plt.grid()

  plt.yscale('log')
  plt.xscale('log')

  plt.xlabel(r'Mesh spacing $\delta x$')
  plt.ylabel("Error norm")

  plt.savefig("norm.pdf")

  #plt.show()
  plt.close()
except:
  # Plotting could fail for any number of reasons, and the actual
  # error raised may depend on, among other things, the current
  # matplotlib backend, so catch everything
  pass

if order > 1.8 and order < 2.2:
  # test for success
  print(" => Test passed")
  exit(0)
else:
  print(" => Test failed")
  exit(1)

